import torch
print(torch.__version__)

x = torch.arange(4, dtype=torch.float)
print(x, x.sum())

A = torch.arange(20, dtype=torch.float).reshape(5, 4)
print(A.shape, A.sum())

print(A)
A_sum_axis0 = A.sum(0)
print(A_sum_axis0, A_sum_axis0.shape)

A_sum_axis1 = A.sum(1)
print(A_sum_axis1, A_sum_axis1.shape)

print(A.mean(), A.sum()/A.numel())

print(A.mean(0), A.sum(0)/A.shape[0])

sum_A = A.sum(axis = 1, keepdims = True)
print(sum_A)
print(A/sum_A)
print(A.cumsum(axis=0))